import os
from typing import List, Dict, Any
from openai import OpenAI


class RAGSystem:
    """RAG 问答系统"""

    def __init__(self):
        self.client = OpenAI(
            api_key=os.getenv("SILICONFLOW_API_KEY"),
            base_url=os.getenv("SILICONFLOW_BASE_URL")
        )
        self.llm_model = os.getenv("LLM_MODEL", "deepseek-ai/DeepSeek-V3")

    def generate_answer(self, question: str, context_docs: List, conversation_history: List[Dict] = None) -> Dict[
        str, Any]:
        """基于检索到的上下文生成答案"""

        # 构建上下文
        context = "\n\n".join([doc.page_content for doc in context_docs])

        # 构建系统提示词
        system_prompt = """你是一个专业的知识库助手。请根据提供的上下文信息回答问题。

要求：
1. 严格基于提供的上下文信息回答
2. 如果上下文没有相关信息，请明确说明"根据现有知识库无法回答该问题"
3. 回答要准确、简洁、有用
4. 不要编造不存在的信息

上下文信息：
{context}
"""
        # 构建消息历史
        messages = []

        # 添加系统提示词
        messages.append({
            "role": "system",
            "content": system_prompt.format(context=context)
        })

        # 添加对话历史
        if conversation_history:
            messages.extend(conversation_history)

        # 添加当前问题
        messages.append({
            "role": "user",
            "content": question
        })

        try:
            # 调用大模型
            response = self.client.chat.completions.create(
                model=self.llm_model,
                messages=messages,
                max_tokens=2000,
                temperature=0.3,  # 较低的温度保证准确性
                stream=False
            )

            answer = response.choices[0].message.content

            return {
                "answer": answer,
                "source_documents": [
                    {
                        "content": doc.page_content[:500] + "..." if len(doc.page_content) > 500 else doc.page_content,
                        "metadata": doc.metadata
                    } for doc in context_docs
                ],
                "tokens_used": response.usage.total_tokens if response.usage else None
            }

        except Exception as e:
            return {
                "error": f"生成答案失败: {str(e)}",
                "answer": None,
                "source_documents": []
            }


# 初始化 RAG 系统
rag_system = RAGSystem()